from __future__ import annotations

import unittest

import numpy as np
from comp_sys import CompLabeledSys, IsNoPBC
from context import dpdata


class TestOrcaSP(unittest.TestCase, CompLabeledSys, IsNoPBC):
    def setUp(self):
        energy_convert = dpdata.unit.EnergyConversion("hartree", "eV").value()
        force_convert = dpdata.unit.ForceConversion(
            "hartree/bohr", "eV/angstrom"
        ).value()

        self.system_1 = dpdata.LabeledSystem("orca/orca.spout", fmt="orca/spout")

        self.system_2 = dpdata.LabeledSystem(
            data={
                "atom_types": np.array(
                    [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 2, 1, 1, 1, 1, 3, 1, 3, 1]
                ),
                "atom_names": ["C", "H", "N", "O"],
                "atom_numbs": [8, 11, 1, 2],
                "coords": np.array(
                    [
                        [
                            [-1.74744e00, 2.17247e00, 3.84400e-02],
                            [-3.05879e00, 1.67227e00, 3.01100e-02],
                            [-3.27420e00, 2.83660e-01, -7.86000e-03],
                            [-2.17997e00, -5.85250e-01, -3.67900e-02],
                            [-6.42570e-01, 1.30482e00, 1.03900e-02],
                            [-8.75230e-01, -8.54100e-02, -2.75900e-02],
                            [-1.58753e00, 3.24402e00, 6.79500e-02],
                            [-6.23100e-02, -7.96870e-01, -5.06600e-02],
                            [-2.34094e00, -1.65578e00, -6.62400e-02],
                            [7.43270e-01, 1.92237e00, 2.35500e-02],
                            [1.91059e00, 9.25720e-01, 4.06000e-03],
                            [8.36210e-01, 2.54508e00, 9.40240e-01],
                            [8.33280e-01, 2.58745e00, -8.63170e-01],
                            [3.18351e00, 1.63882e00, 3.70200e-02],
                            [1.86317e00, 2.99040e-01, -9.14520e-01],
                            [1.84793e00, 2.69100e-01, 8.99400e-01],
                            [3.28398e00, 2.20442e00, -8.37230e-01],
                            [3.95753e00, 9.35910e-01, 5.24200e-02],
                            [-4.10991e00, 2.51820e00, 5.88000e-02],
                            [-3.99914e00, 3.47933e00, 8.64100e-02],
                            [-4.52084e00, -2.33960e-01, -1.68400e-02],
                            [-5.31660e00, 3.15930e-01, 2.60000e-03],
                        ]
                    ]
                ),
                "energies": np.array([-513.113388868587]) * energy_convert,
                "forces": np.array(
                    [
                        [
                            [-5.7786180e-03, -2.4701072e-02, -6.2814900e-04],
                            [3.2387534e-02, 2.0888587e-02, 4.9131600e-04],
                            [-1.3022767e-02, 1.3820567e-02, 4.5543300e-04],
                            [-4.4279100e-04, 1.4037682e-02, 4.7424800e-04],
                            [-6.3908530e-03, -8.4241470e-03, -1.2989690e-03],
                            [-8.9298450e-03, 1.6404774e-02, 5.9260500e-04],
                            [-1.0737810e-03, -1.5698400e-04, 4.1875000e-05],
                            [4.6807600e-04, 2.2149790e-03, -4.0880000e-06],
                            [-3.9808540e-03, 2.3696040e-03, 1.0808900e-04],
                            [5.0405390e-03, 1.0389430e-03, 2.3300390e-03],
                            [3.1264600e-03, 2.7682000e-04, -3.8873520e-03],
                            [1.8491730e-03, -4.7411320e-03, -1.3538358e-02],
                            [3.7527600e-04, -4.2966570e-03, 1.1656289e-02],
                            [2.0499798e-02, -7.3734830e-03, -2.0336553e-02],
                            [-6.7151250e-03, 1.3060670e-03, 1.2263601e-02],
                            [1.6883400e-03, 6.5851660e-03, -1.2895903e-02],
                            [5.3261000e-04, -2.0068781e-02, 2.8391439e-02],
                            [-2.2997109e-02, 2.6866202e-02, -3.1128700e-03],
                            [-2.7960063e-02, 6.0005960e-03, 2.3320500e-04],
                            [1.5089874e-02, -3.2118582e-02, -9.5217200e-04],
                            [-1.9753810e-02, -1.1993684e-02, -2.8813700e-04],
                            [3.5987936e-02, 2.0645360e-03, -9.5588000e-05],
                        ]
                    ]
                )
                * force_convert,
                "cells": np.zeros((1, 3, 3)),
                "orig": np.zeros(3),
                "nopbc": True,
            }
        )
        self.places = 6
        self.e_places = 9
        self.f_places = 9
        self.v_places = 6
